import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
df = pd.read_csv("../datas/bayes_wangzhe.txt",header=None)
X = df[1]
Y = df[0]
tfCoder = TfidfVectorizer(token_pattern="[a-zA-Z|\u4e00-\u9fa5]+")
X = tfCoder.fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2, random_state=42)
from sklearn.naive_bayes import MultinomialNB, GaussianNB, BernoulliNB
model = MultinomialNB()
model.fit(X_train,y_train)
# print(model.predict(X_train))
# print(tfCoder.get_feature_names())
# print(X_train.toarray())
# print(y_train.values)
a = ["残血 的 安琪拉 打 不 过 鲁班","这 一波 大龙 别 再 被 抢 了",
     "你 在 石头 那 不要 动，我 去 买 几个 橘子"]
# print(tfCoder.transform(a).todense())
print(model.predict(tfCoder.transform(a)))
